// jnijvpx8.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include <iostream>
#include <jni.h>
#include "vpx_image.h"
#include "vpx_codec.h"
#include "vp8.h"
#include "inttypes.h"
#include "stdint.h"
#include "vpx_codec.h"
#include "vp8.h"
#include "vpx_decoder.h"
#include "vpx_encoder.h"
#include "vp8cx.h"

#define interfac (_vpx_codec_vp8_cx())
#define interfacDecoder (_vpx_codec_vp8_dx())
#define fourcc 0x30385056
#define IVF_FILE_HDR_SZ (32)
#define IVF_FRAME_HDR_SZ (12)

extern const char*(*_vpx_codec_build_config)(void);
extern vpx_codec_err_t(*_vpx_codec_control_)(vpx_codec_ctx_t  *ctx, int ctrl_id,...);
extern vpx_codec_err_t(*_vpx_codec_dec_init_ver)(vpx_codec_ctx_t*ctx,vpx_codec_iface_t*iface,vpx_codec_dec_cfg_t*cfg,vpx_codec_flags_t flags,int ver);
extern vpx_codec_err_t(*_vpx_codec_decode)(vpx_codec_ctx_t*ctx,const uint8_t*data,unsigned int data_sz,void*user_priv,long deadline);
extern vpx_codec_err_t(*_vpx_codec_destroy)(vpx_codec_ctx_t *ctx);
extern vpx_codec_err_t(*_vpx_codec_enc_config_default)(vpx_codec_iface_t*iface,vpx_codec_enc_cfg_t*cfg,unsigned int usage);
extern vpx_codec_err_t(*_vpx_codec_enc_config_set)(vpx_codec_ctx_t*ctx,const vpx_codec_enc_cfg_t*cfg);
extern vpx_codec_err_t(*_vpx_codec_enc_init_ver)(vpx_codec_ctx_t*ctx,vpx_codec_iface_t*iface,vpx_codec_enc_cfg_t*cfg,vpx_codec_flags_t flags,int ver);
extern vpx_codec_err_t(*_vpx_codec_encode)(vpx_codec_ctx_t*ctx,const vpx_image_t*img,vpx_codec_pts_t pts,unsigned long duration,vpx_enc_frame_flags_t flags,unsigned long deadline);
extern const char*(*_vpx_codec_err_to_string)(vpx_codec_err_t  err);
extern const char*(*_vpx_codec_error)(vpx_codec_ctx_t*ctx);
extern const char*(*_vpx_codec_error_detail)(vpx_codec_ctx_t*ctx);
extern vpx_codec_caps_t(*_vpx_codec_get_caps)(vpx_codec_iface_t *iface);
extern const vpx_codec_cx_pkt_t*(*_vpx_codec_get_cx_data)(vpx_codec_ctx_t*ctx,vpx_codec_iter_t*iter);
extern vpx_image_t*(*_vpx_codec_get_frame)(vpx_codec_ctx_t  *ctx,vpx_codec_iter_t *iter);
extern vpx_fixed_buf_t*(*_vpx_codec_get_global_headers)(vpx_codec_ctx_t*ctx);
extern vpx_codec_err_t(*_vpx_codec_get_mem_map)(vpx_codec_ctx_t*ctx,vpx_codec_mmap_t*mmap,vpx_codec_iter_t*iter);
extern const vpx_image_t*(*_vpx_codec_get_preview_frame)(vpx_codec_ctx_t*ctx);
extern vpx_codec_err_t(*_vpx_codec_get_stream_info)(vpx_codec_ctx_t*ctx,vpx_codec_stream_info_t*si);
extern const char*(*_vpx_codec_iface_name)(vpx_codec_iface_t *iface);
extern vpx_codec_err_t(*_vpx_codec_peek_stream_info)(vpx_codec_iface_t*iface,const uint8_t*data,unsigned int data_sz,vpx_codec_stream_info_t*si);
extern vpx_codec_err_t(*_vpx_codec_register_put_frame_cb)(vpx_codec_ctx_t*ctx,vpx_codec_put_frame_cb_fn_t cb,void*user_priv);
extern vpx_codec_err_t(*_vpx_codec_register_put_slice_cb)(vpx_codec_ctx_t*ctx,vpx_codec_put_slice_cb_fn_t cb,void*user_priv);
extern vpx_codec_err_t(*_vpx_codec_set_cx_data_buf)(vpx_codec_ctx_t*ctx,const vpx_fixed_buf_t*buf,unsigned int pad_before,unsigned int pad_after);
extern vpx_codec_err_t(*_vpx_codec_set_mem_map)(vpx_codec_ctx_t*ctx,vpx_codec_mmap_t*mmaps,unsigned int num_maps);
extern int(*_vpx_codec_version)(void);
extern const char*(*_vpx_codec_version_extra_str)(void);
extern const char*(*_vpx_codec_version_str)(void);
extern vpx_codec_iface_t* (*_vpx_codec_vp8_cx)(void);
extern vpx_codec_iface_t* (*_vpx_codec_vp8_dx)(void);
extern vpx_image_t*(*_vpx_img_alloc)(vpx_image_t  *img,vpx_img_fmt_t fmt,unsigned int d_w,unsigned int d_h,unsigned int align);
extern void(*_vpx_img_flip)(vpx_image_t*img);
extern void(*_vpx_img_free)(vpx_image_t*img);
extern int(*_vpx_img_set_rect)(vpx_image_t*img,unsigned int x,unsigned int y,unsigned int w,unsigned int h);
extern vpx_image_t*(*_vpx_img_wrap)(vpx_image_t*img,vpx_img_fmt_t fmt,unsigned int d_w,unsigned int d_h,unsigned int align,unsigned char*img_data);
extern void loadLib(const wchar_t*libpath);

jfieldID decoderCodecPointerField;
jfieldID decoderWidthField;
jfieldID decoderHeightField;

jfieldID codecPointerField;
jfieldID rawVpxImageField;
jfieldID flagsField;
jfieldID widthField;
jfieldID heightField;
jfieldID bitrateField;
jfieldID maxDistField;
jfieldID frameCntField;
jfieldID cfgField;
jmethodID writeMethod;

static inline unsigned char clampInt(int value){
	if(value<0)return 0;
	if(value>255)return 255;
	return (unsigned char)value;
}

static void mem_put_le16(char *mem, unsigned int val) {
	mem[0] = val;
	mem[1] = val>>8;
}

static void mem_put_le32(char *mem, unsigned int val) {
	mem[0] = val;
	mem[1] = val>>8;
	mem[2] = val>>16;
	mem[3] = val>>24;
}

#define rgbtoy(b, g, r, y) \
	y=(unsigned char)(((int)(30*r) + (int)(59*g) + (int)(11*b))/100)


//rgbtoy(b, g, r, y); \

#define rgbtoyuv(b, g, r, y, u, v) \
	y=(unsigned char)clampInt(((int)(30*r) + (int)(59*g) + (int)(11*b))/100);\
	u=(unsigned char)(((int)(-17*r) - (int)(33*g) + (int)(50*b)+12800)/100); \
	v=(unsigned char)(((int)(50*r) - (int)(42*g) - (int)(8*b)+12800)/100)

int sizeOfUv12Frame(int width, int height){
	return width*height*3/2;
}

static void RGBtoYUV420PSameSize (const unsigned char * rgb,
	unsigned char * yuv,
	unsigned rgbIncrement,
	unsigned char flip,
	int srcFrameWidth, int srcFrameHeight) {
		unsigned int planeSize;
		unsigned int halfWidth;

		unsigned char * yplane;
		unsigned char * uplane;
		unsigned char * vplane;
		const unsigned char * rgbIndex;

		int x, y;
		unsigned char * yline;
		unsigned char * uline;
		unsigned char * vline;

		planeSize = srcFrameWidth * srcFrameHeight;
		halfWidth = srcFrameWidth >> 1;

		// get pointers to the data
		yplane = yuv;
		uplane = yuv + planeSize;
		vplane = yuv + planeSize + (planeSize >> 2);
		rgbIndex = rgb;

		for (y = 0; y < srcFrameHeight; y++)
		{
			yline = yplane + (y * srcFrameWidth);
			uline = uplane + ((y >> 1) * halfWidth);
			vline = vplane + ((y >> 1) * halfWidth);

			if (flip)
				rgbIndex = rgb + (srcFrameWidth*(srcFrameHeight-1-y)*rgbIncrement);

			for (x = 0; x < (int) srcFrameWidth; x+=2)
			{
				rgbtoyuv(rgbIndex[0], rgbIndex[1], rgbIndex[2], *yline, *uline, *vline);
				rgbIndex += rgbIncrement;
				yline++;
				rgbtoyuv(rgbIndex[0], rgbIndex[1], rgbIndex[2], *yline, *uline, *vline);
				rgbIndex += rgbIncrement;
				yline++;
				uline++;
				vline++;
			}
		}
		//printf("yline=%d uline=%d vline=%d rgb=%d\n",yline-yuv,uline-yuv, vline-yuv, rgbIndex-rgb);
}

static inline unsigned char clamp(float value){
	if(value<0)return 0;
	if(value>255)return 255;
	return (unsigned char)value;
}

static inline unsigned char clampi(int value){
	if(value<0)return 0;
	if(value>255)return 255;
	return (unsigned char)value;
}

#define yuvtorgb(b, g, r, y, u, v) \
	r=clamp((float)y + 1.13983f * ((float)v - 128)); \
	g=clamp((float)y - 0.39465f * ((float)u - 128) - 0.58060f * ((float)v - 128)); \
	b=clamp((float)y + 2.03211f * ((float)u - 128))



static void YUV420PtoRgbSameSize (unsigned char * rgb,
	unsigned char * yuv,
	unsigned rgbIncrement,
	unsigned char flip,
	int srcFrameWidth, int srcFrameHeight) {
		unsigned int planeSize;
		unsigned int halfWidth;

		unsigned char * yplane;
		unsigned char * uplane;
		unsigned char * vplane;
		unsigned char * rgbIndex;

		int x, y;
		unsigned char * yline;
		unsigned char * uline;
		unsigned char * vline;

		planeSize = srcFrameWidth * srcFrameHeight;
		halfWidth = srcFrameWidth >> 1;

		// get pointers to the data
		yplane = yuv;
		uplane = yuv + planeSize;
		vplane = yuv + planeSize + (planeSize >> 2);
		rgbIndex = rgb;

		for (y = 0; y < srcFrameHeight; y++)
		{
			yline = yplane + (y * srcFrameWidth);
			uline = uplane + ((y >> 1) * halfWidth);
			vline = vplane + ((y >> 1) * halfWidth);

			if (flip)
				rgbIndex = rgb + (srcFrameWidth*(srcFrameHeight-1-y)*rgbIncrement);

			for (x = 0; x < (int) srcFrameWidth; x+=2)
			{
				yuvtorgb(rgbIndex[0], rgbIndex[1], rgbIndex[2], *yline, *uline, *vline);
				rgbIndex += rgbIncrement;
				yline++;
				yuvtorgb(rgbIndex[0], rgbIndex[1], rgbIndex[2], *yline, *uline, *vline);
				rgbIndex += rgbIncrement;
				yline++;
				uline++;
				vline++;
			}
		}
		//printf("yline=%d uline=%d vline=%d rgb=%d\n",yline-yuv,uline-yuv, vline-yuv, rgbIndex-rgb);
}


extern "C" {

	JNIEXPORT jbyteArray JNICALL Java_com_vpx_JniVpx_convertImageToYuv(JNIEnv *env, jclass obj, jlong address, jint width, jint height){
		int size=sizeOfUv12Frame(width, height);
		unsigned char*resultArray=new unsigned char[size];
		RGBtoYUV420PSameSize((const unsigned char*)address,resultArray,3,0,width, height);
		jbyteArray ba=env->NewByteArray(size);
		env->SetByteArrayRegion(ba,0,size,(const jbyte*)resultArray);
		//delete [] resultArray;
		return ba;
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_convertYuvToImage(JNIEnv *env, jclass obj, jlong address, jint width, jint height, jbyteArray yuvArray){
		jbyte*yuvData=env->GetByteArrayElements(yuvArray, 0);
		YUV420PtoRgbSameSize((unsigned char*)address,(unsigned char*)yuvData,3,0,width, height);
		env->ReleaseByteArrayElements(yuvArray,yuvData, 0);
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_loadVpxLib(JNIEnv *env, jclass obj,jstring libPath){
		std::wstring value;
		const jchar *raw = env->GetStringChars(libPath, 0);
		jsize len = env->GetStringLength(libPath);
		const jchar *temp = raw;
		value.assign(raw, raw + len);
		env->ReleaseStringChars(libPath, raw);
		loadLib(value.c_str());
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_createDecoder(JNIEnv *env, jclass obj,jobject javaDecoder){
		jclass cls = env->FindClass("com/vpx/VpxDecoder");
		decoderCodecPointerField = env->GetFieldID(cls, "codecPointer", "J");
		decoderWidthField = env->GetFieldID(cls, "width", "I");
		decoderHeightField = env->GetFieldID(cls, "height", "I");
		vpx_codec_ctx_t*codec=new vpx_codec_ctx_t();		
		if(_vpx_codec_dec_init_ver(codec, interfacDecoder, NULL, 0, VPX_DECODER_ABI_VERSION)){
			char*exBuffer=new char[65000];
			sprintf(exBuffer, "Failed to initialize decoder %s", _vpx_codec_error(codec));
			env->ThrowNew(env->FindClass("java/lang/Exception"), exBuffer);
			return;
		}

		env->SetLongField(javaDecoder, decoderCodecPointerField,(jlong)codec);
	}

	JNIEXPORT jboolean JNICALL Java_com_vpx_JniVpx_decodeFrame(JNIEnv *env, jclass obj,jobject javaDecoder, jbyteArray frameByteArray, jlong address){
		vpx_codec_ctx_t*codec=(vpx_codec_ctx_t*)env->GetLongField(javaDecoder,decoderCodecPointerField);
		jint frameWidth=env->GetIntField(javaDecoder, decoderWidthField);
		jint frameHeight=env->GetIntField(javaDecoder, decoderHeightField);
		jsize frameSize=env->GetArrayLength(frameByteArray);
		jbyte*frameData=env->GetByteArrayElements(frameByteArray, 0);
		if(_vpx_codec_decode(codec, (const uint8_t*)frameData, frameSize, NULL, 1)){
			char*exBuffer=new char[65000];
			sprintf(exBuffer, "Failed to decode frame. %s", _vpx_codec_error(codec));
			env->ThrowNew(env->FindClass("java/lang/Exception"), exBuffer);
			return false;
		}

		vpx_image_t*img =NULL;
		vpx_codec_iter_t  iter = NULL;
		vpx_image_t*tImage =NULL;
		while((img = _vpx_codec_get_frame(codec, &iter))) {  
			if(img!=NULL){
				tImage=img;
			}
		}

		if(tImage==NULL){
			return false;
		}
		if(tImage->d_w!=frameWidth || tImage->d_h!=frameHeight){
			char*exBuffer=new char[65000];
			sprintf(exBuffer, "Decoded frame size %d:%d not equals to expected frame size %d:%d", tImage->w, tImage->h, frameWidth, frameHeight);
			env->ThrowNew(env->FindClass("java/lang/Exception"), exBuffer);
			return false;
		}
		unsigned char*addr=(unsigned char*)address;
		int height = tImage->d_h, width = tImage->d_w;
        unsigned char *ybuf = tImage->planes[0];
        unsigned char *ubuf = tImage->planes[1];
        unsigned char *vbuf = tImage->planes[2];
		int rowLength=width*3;
		for(int y = 0; y < height; y++){
            unsigned char *yy = ybuf, *uu = ubuf, *vv = vbuf;
			int pos = y*rowLength;
			for(int x = 0; x < width; x++){
                int y0 = ((*yy) - 16) * 100000, u0 = (int)(*uu) - 128, v0 = (int)(*vv) - 128;
                int r = (y0 + 140200 * v0) / 100000;
                int g = (y0 - 34414 * u0 - 71414 * v0) / 100000;
                int b = (y0 + 177200 * u0) / 100000;

                
                addr[pos] = clampi(b);
                addr[pos + 1] = clampi(g);
                addr[pos + 2] = clampi(r);
                
				pos+=3;
                yy++;
                if((x%2)==0){
				uu++;
                vv++;
				}
            }
            ybuf += tImage->stride[0];
            if((y%2)==0){
				ubuf += tImage->stride[1];
				vbuf += tImage->stride[2];
			}
        }
		
		return true;
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_createEncoder(JNIEnv *env, jclass obj,jobject javaCodec){
		jclass cls = env->FindClass("com/vpx/VpxEncoder");
		codecPointerField = env->GetFieldID(cls, "codecPointer", "J");
		rawVpxImageField = env->GetFieldID(cls, "rawVpxImage", "J");
		flagsField = env->GetFieldID(cls, "flags", "I");
		widthField = env->GetFieldID(cls, "width", "I");
		heightField = env->GetFieldID(cls, "height", "I");
		bitrateField = env->GetFieldID(cls, "bitrate", "I");
		maxDistField = env->GetFieldID(cls, "max_dist", "I");
		frameCntField = env->GetFieldID(cls, "frame_cnt", "I");
		cfgField = env->GetFieldID(cls, "cfg", "J");
		jclass byteStreamClass=env->FindClass("java/io/ByteArrayOutputStream");
		writeMethod=env->GetMethodID(byteStreamClass,"write","([BII)V");

		int width=env->GetIntField(javaCodec, widthField);
		int height=env->GetIntField(javaCodec, heightField);
		int bitrate=env->GetIntField(javaCodec, bitrateField);
		int maxDist=env->GetIntField(javaCodec, maxDistField);

		vpx_codec_ctx_t*codec=new vpx_codec_ctx_t();
		vpx_image_t*rawImage=new vpx_image_t();
		if(!_vpx_img_alloc(rawImage, VPX_IMG_FMT_I420, width, height, 1)){
			char*exBuffer=new char[65000];
			sprintf(exBuffer, "Failed to allocate image %d %d. %s", width, height, _vpx_codec_error(codec));
			env->ThrowNew(env->FindClass("java/lang/Exception"), exBuffer);
		}

		vpx_codec_enc_cfg_t*cfg = new vpx_codec_enc_cfg_t();
		_vpx_codec_enc_config_default(interfac,cfg, 0);
		cfg->rc_target_bitrate = bitrate;
		cfg->g_w = width;
		cfg->g_h = height;    
		cfg->kf_max_dist=maxDist;

		if(_vpx_codec_enc_init_ver(codec, interfac, cfg, 0,VPX_ENCODER_ABI_VERSION)){
			env->ThrowNew(env->FindClass("java/lang/Exception"), "Failed to initialize encoder");
		}

		env->SetLongField(javaCodec, codecPointerField, (jlong)codec);
		env->SetLongField(javaCodec, rawVpxImageField, (jlong)rawImage);
		env->SetLongField(javaCodec, cfgField, (jlong)cfg);
	}

	void writeArrayToByteArrayStream(JNIEnv*env, jobject byteArrayOutputStream, void*bytes, int count){
		jbyteArray bArray = env->NewByteArray(count);
		env->SetByteArrayRegion(bArray,0,count,(const jbyte*)bytes);
		env->CallVoidMethod(byteArrayOutputStream,writeMethod,bArray, 0, count);
		//env->DeleteLocalRef(bArray);
	}

	static void write_ivf_frame_header(JNIEnv*env, jobject byteArrayOutputStream,const vpx_codec_cx_pkt_t *pkt){
		char header[12];
		vpx_codec_pts_t pts;
		if(pkt->kind != VPX_CODEC_CX_FRAME_PKT){
			return;
		}

		pts = pkt->data.frame.pts;
		mem_put_le32(header, pkt->data.frame.sz);
		mem_put_le32(header+4, pts&0xFFFFFFFF);
		mem_put_le32(header+8, pts >> 32);
		writeArrayToByteArrayStream(env, byteArrayOutputStream,header, 12);
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_encodeFrame(JNIEnv *env, jclass obj, jlong address,jobject javaCodec, jobject byteArrayOutputStream){
		int width=env->GetIntField(javaCodec,widthField);
		int height=env->GetIntField(javaCodec,heightField);
		int frame_cnt=env->GetIntField(javaCodec, frameCntField);
		int flags=env->GetIntField(javaCodec, flagsField);

		vpx_codec_ctx_t*codec=(vpx_codec_ctx_t*)env->GetLongField(javaCodec, codecPointerField);
		vpx_image_t*raw=(vpx_image_t*)env->GetLongField(javaCodec, rawVpxImageField);
		vpx_codec_iter_t iter = NULL;
		const vpx_codec_cx_pkt_t *pkt;
		const unsigned char*addr=(const unsigned char*)address;
		RGBtoYUV420PSameSize(addr,raw->planes[0],3,0,width, height);

		if(_vpx_codec_encode(codec, raw, frame_cnt,1, flags, VPX_DL_REALTIME)){
			env->ThrowNew(env->FindClass("java/lang/Exception"), "Failed to encode frame");	
		}
		int count=0;
		while( (pkt = _vpx_codec_get_cx_data(codec, &iter)) ) {
			switch(pkt->kind) {
			case VPX_CODEC_CX_FRAME_PKT:
				//write_ivf_frame_header(env, byteArrayOutputStream,pkt);
				writeArrayToByteArrayStream(env, byteArrayOutputStream, pkt->data.frame.buf, pkt->data.frame.sz);
				break;
			default:
				break;
			}

			count++;
		}

		if(count>1){
			printf("Count %d", count);
		}

		fflush(stdout);
		frame_cnt++;
		env->SetIntField(javaCodec, frameCntField, frame_cnt);
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_destroyEncoder(JNIEnv *env, jclass obj, jobject javaCodec){
		vpx_codec_ctx_t*codec=(vpx_codec_ctx_t*)env->GetLongField(javaCodec, codecPointerField);
		vpx_image_t*raw=(vpx_image_t*)env->GetLongField(javaCodec, rawVpxImageField);
		vpx_codec_enc_cfg_t*cfg=(vpx_codec_enc_cfg_t*)env->GetLongField(javaCodec, cfgField);

		_vpx_img_free(raw);
		if(_vpx_codec_destroy(codec)){
			env->ThrowNew(env->FindClass("java/lang/Exception"), "Failed to destroy encoder");
		}
		delete codec;
		delete raw;
	}

	JNIEXPORT void JNICALL Java_com_vpx_JniVpx_destroyDecoder(JNIEnv *env, jclass obj, jobject javaDecoder){
		vpx_codec_ctx_t*codec=(vpx_codec_ctx_t*)env->GetLongField(javaDecoder, decoderCodecPointerField);
		if(_vpx_codec_destroy(codec)){
			env->ThrowNew(env->FindClass("java/lang/Exception"), "Failed to destroy decoder");
		}
		delete codec;
	}

	/*JNIEXPORT void JNICALL Java_com_vpx_JniVpx_close(JNIEnv *env, jclass obj){
	_vpx_img_free(&raw);
	if(_vpx_codec_destroy(&codec)){
	die(env,"Failed to destroy codec");
	}

	//Try to rewrite the file header with the actual frame count 
	if(!fseek(outfile, 0, SEEK_SET)){
	write_ivf_file_header(outfile, &cfg, frame_cnt-1);
	}

	fclose(outfile);
	}
	*/
}

